python equivalent of matlab find